package cn.iocoder.yudao.module.engine.model.aichat;

import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.locks.ReentrantLock;

@Slf4j
public class AIChatQueue {
    private final LinkedList<AIChatFuture> queue;
    private final ReentrantLock lock;

    public AIChatQueue() {
        queue = new LinkedList<>();
        lock = new ReentrantLock();
    }

    public List<AIChatFuture> add(AIChatFuture future) {
        lock.lock();
        List<AIChatFuture> removedMessages = new ArrayList<>();
        try {
            String newChatId = future.getChatId();
            boolean chatIdExists = false;
            int insertIndex = 0;

            // 检查 chatId 是否已存在于队列中
            for (int i = queue.size() - 1; i >= 0; i--) {
                AIChatFuture currentMessage = queue.get(i);
                if (currentMessage.getChatId().equals(newChatId)) {
                    chatIdExists = true;
                    insertIndex = i + 1;
                    break;
                }
            }

            if (chatIdExists) {
                // 如果 chatId 已存在，插入到相同 chatId 的下面
                queue.add(insertIndex, future);
                log.info("chatId已存在，插入到相同chatId的下面, newChatId = {}", newChatId);
            } else {
                int differentChatIds = getDifferentChatIdCount();
                if (differentChatIds < 2) {
                    // 如果不同的 chatId 数量不超过 2 个，直接添加消息
                    queue.add(future);
                    log.info("不同的chatId数量不超过2个，直接添加消息, newChatId = {}", newChatId);
                } else {
                    // 如果不同的 chatId 数量超过 2 个，清空最新的 chatId 相关的消息
                    final String latestChatId = queue.getLast().getChatId();
                    for (int i = queue.size() - 1; i >= 0; i--) {
                        AIChatFuture msg = queue.get(i);
                        if (msg.getChatId().equals(latestChatId)) {
                            removedMessages.add(queue.remove(i));
                        }
                    }
                    if (!removedMessages.isEmpty()) {
                        log.info("chatId超过2个，移除了最新的chatId相关的消息, latestChatId = {}", latestChatId);
                    }
                    // 添加新消息
                    queue.add(future);
                    log.info("chatId超过2个，添加新消息, newChatId = {}", newChatId);
                }
            }
        } finally {
            lock.unlock();
        }
        return removedMessages;
    }

    public boolean remove(AIChatFuture future) {
        lock.lock();
        try {
            if (queue.isEmpty()) {
                return false;
            }
            return queue.remove(future);
        } finally {
            lock.unlock();
        }
    }

    public AIChatFuture remove() {
        lock.lock();
        try {
            if (queue.isEmpty()) {
                return null;
            }
            return queue.removeFirst();
        } finally {
            lock.unlock();
        }
    }

    public void clear() {
        lock.lock();
        try {
            queue.clear();
        } finally {
            lock.unlock();
        }
    }

    private int getDifferentChatIdCount() {
        int count = 0;
        String firstChatId = null;
        String secondChatId = null;
        for (AIChatFuture message : queue) {
            String chatId = message.getChatId();
            if (firstChatId == null) {
                firstChatId = chatId;
                count++;
            } else if (!chatId.equals(firstChatId) && (secondChatId == null || !chatId.equals(secondChatId))) {
                if (secondChatId == null) {
                    secondChatId = chatId;
                    count++;
                }
            }
        }
        return count;
    }

    public List<AIChatFuture> getQueue() {
        lock.lock();
        try {
            return new ArrayList<>(queue);
        } finally {
            lock.unlock();
        }
    }

    public int size() {
        lock.lock();
        try {
            return queue.size();
        } finally {
            lock.unlock();
        }
    }

    public static void main(String[] args) {
        AIChatQueue audioQueue = new AIChatQueue();

        audioQueue.add(new AIChatFuture("a", "a"));
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        audioQueue.add(new AIChatFuture("b", "b"));
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        audioQueue.add(new AIChatFuture("a", "a"));
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        audioQueue.add(new AIChatFuture("b", "b"));
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        List<AIChatFuture> removed = audioQueue.add(new AIChatFuture("c", "c"));
        System.out.println("Removed messages: " + removed.stream().map(v -> v.getChatId()).toList());
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        audioQueue.add(new AIChatFuture("c", "c"));
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        audioQueue.remove();
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());

        audioQueue.add(new AIChatFuture("bgm-a", "bgm-a"));
        System.out.println(audioQueue.getQueue().stream().map(v -> v.getChatId()).toList());
    }
}